#!/usr/bin/env python
# -*- coding: utf-8 -*-

from PyQt4.QtCore import *
from PyQt4.QtXml import *

def intFromQStr(qstr):
    i, ok = qstr.toInt()
    if not ok:
        raise ValueError, unicode(qstr)
    return i

def sprint(x):
  print x

class Verticle:

  def __init__(self, data):
    self.data = data
    self.connectedVerticles = set()
    self.visited = False
    self.index = -1

  def __str__(self):
    return "<Verticle data=\"%s\" index=\"%d\"/>\n" % (self.data, self.index)

  def execute(self, action):
    action(self.data)

  def addConnection(self, other):
    if (other != self):
      self.connectedVerticles.add(other)
  
  def delConnection(self, other):
    if (other != self and other in self.connectedVerticles):
      self.connectedVerticles.remove(other)

class Graph:

  """ Unit test code:

    >>> graph = Graph()
    >>> a = graph.addVerticle("A")
    >>> b = graph.addVerticle("B")
    >>> c = graph.addVerticle("C")
    >>> d = graph.addVerticle("D")
    >>> e = graph.addVerticle("E")
    >>> f = graph.addVerticle("F")
    >>> g = graph.addVerticle("G")
    >>> graph.addEdge(a, b)
    >>> graph.addEdge(b, c)
    >>> graph.addEdge(c, d)
    >>> graph.addEdge(d, e)
    >>> graph.addEdge(e, f)
    >>> graph.saveToFile("./test.xml")
    (True, 'Saved graph into ./test.xml')
    >>> graph2 = Graph("./test.xml")
    >>> graph.dsf(a, lambda x: sprint(x))
    A
    B
    C
    D
    E
    F
    >>> graph2.dsf(a, lambda x: sprint(x))
    A
    B
    C
    D
    E
    F
    """

  verticlesCount = 0

  def __init__(self, filename=None):
    self.verticles = dict()
    if not filename is None:
      self.loadFromFile(filename)

  def getVerticle(self, index):
    return self.verticles.get(index)

  def getVerticles(self):
    return self.verticles.keys()

  def addVerticle(self, data):
    verticle = Verticle(data)
    self.verticles[Graph.verticlesCount] = verticle
    verticle.index = Graph.verticlesCount
    Graph.verticlesCount += 1
    return Graph.verticlesCount - 1

  def delVerticle(self, index):
    vertToDel = self.getVerticle(index)
    if vertToDel is not None:
      for verticle in vertToDel.connectedVerticles:
        verticle.delConnection(vertToDel)
      del self.verticles[index]
  
  def addEdge(self, index_1, index_2):
    verticle_1 = self.getVerticle(index_1)
    verticle_2 = self.getVerticle(index_2)
    if (verticle_1 is not None and verticle_2 is not None):
      verticle_1.addConnection(verticle_2)
      verticle_2.addConnection(verticle_1)

  def delEdge(self, index_1, index_2):
    verticle_1 = self.getVerticle(index_1)
    verticle_2 = self.getVerticle(index_2)
    if (verticle_1 is not None and verticle_2 is not None):
      verticle_1.delConnection(verticle_2)
      verticle_2.delConnection(verticle_1)

  def getConnectedVerticles(self, index):
    verticle = self.getVerticle(index)
    if verticle is None:
      return None
    connVerticles = []
    for connVerticle in verticle.connectedVerticles:
      connVerticles.append(connVerticle.index)
    return connVerticles

  def dsf(self, index, action):
    rootVerticle = self.getVerticle(index)
    self.doDsf(rootVerticle, action)
    for verticle in self.verticles.values():
      verticle.visited = False

  def doDsf(self, rootVerticle, action):
    if (rootVerticle is not None and not rootVerticle.visited):
      rootVerticle.execute(action)
      rootVerticle.visited = True
      for verticle in rootVerticle.connectedVerticles:
        self.doDsf(verticle, action)

  def saveToFile(self, filename):
    error = None
    fh = None
    try:
      fh = QFile(filename)
      if not fh.open(QIODevice.WriteOnly):
        raise IOError, unicode(fh.errorString())
      stream = QTextStream(fh)
      stream << "<Graph>\n"
      for verticle in self.verticles.values():
        stream << "\t%s" % str(verticle)
      for verticle in self.verticles.values():
        for connectedVerticle in verticle.connectedVerticles:
          stream << "\t<Edge first=\"%d\" second=\"%d\"/>\n" % (verticle.index, connectedVerticle.index)
      stream << "</Graph>\n"
    except (IOError, OSError), e:
      error = "Failed to save: %s" % e
    finally:
      if fh is not None:
        fh.close()
      if error is not None:
        return False, error
      return True, "Saved graph into %s" % filename

  def loadFromFile(self, filename): 
    dom = QDomDocument()
    error = None
    fh = None
    try:
        fh = QFile(filename)
        if not fh.open(QIODevice.ReadOnly):
          raise IOError, unicode(fh.errorString())
        if not dom.setContent(fh):
          raise ValueError, "could not parse XML"
    except (IOError, OSError, ValueError), e:
        error = "Failed to import: %s" % e
    finally:
        if fh is not None:
          fh.close()
        if error is not None:
          return False, error
    try:    
        self.populateFromDOM(dom)
    except ValueError, e:
        self.verticles.clear()
        return False, e
    return True, "Loaded from file %s" % filename

  def populateFromDOM(self, dom):
    self.verticles.clear()
    root = dom.documentElement()
    if root.tagName() != "Graph":
        raise ValueError, "not a Aims XML file"
    node = root.firstChild()
    while not node.isNull():
        element = node.toElement()
        if element.tagName() == "Verticle":
          self.verticles[intFromQStr(element.attribute("index"))] = Verticle(unicode(element.attribute("data")))
        if element.tagName() == "Edge":
          self.addEdge(intFromQStr(element.attribute("first")), intFromQStr(element.attribute("second")))
        node = node.nextSibling()

if __name__ == "__main__":
  import doctest
  doctest.testmod()